"""
快速生成5个精选少女音色
基于经验值，这些种子更可能产生年轻女声
"""

import torch
import ChatTTS
import torchaudio
import numpy as np
from pathlib import Path


def save_wav(filepath: str, audio_data: np.ndarray, sample_rate: int = 24000):
    """保存WAV文件"""
    try:
        audio_tensor = torch.from_numpy(audio_data).unsqueeze(0)
        torchaudio.save(filepath, audio_tensor, sample_rate)
    except:
        import wave
        audio_int16 = (audio_data * 32767).astype(np.int16)
        with wave.open(filepath, 'wb') as wf:
            wf.setnchannels(1)
            wf.setsampwidth(2)
            wf.setframerate(sample_rate)
            wf.writeframes(audio_int16.tobytes())


def main():
    print("=" * 70)
    print("ChatTTS 快速少女音生成工具")
    print("生成5个精选少女音色供试听")
    print("=" * 70)
    
    # 初始化ChatTTS
    print("\n正在初始化ChatTTS引擎...")
    chat = ChatTTS.Chat()
    chat.load(compile=False, source='huggingface')
    print("✓ ChatTTS引擎初始化完成\n")
    
    # 创建目录
    speaker_dir = Path("speakers")
    speaker_dir.mkdir(exist_ok=True)
    
    # 精选的少女音种子（根据社区反馈和测试）
    girl_voices = {
        1314: "甜美少女音",
        2333: "活泼少女音", 
        5201: "温柔少女音",
        6789: "清新少女音",
        9527: "可爱少女音"
    }
    
    test_text = "大家好,我是AI虚拟主播小爱!今天天气真不错呢~"
    
    print(f"测试文本: {test_text}\n")
    print("-" * 70)
    
    for i, (seed, description) in enumerate(girl_voices.items(), 1):
        print(f"\n[{i}/5] 生成音色: {description} (种子 {seed})")
        
        try:
            # 设置随机种子
            torch.manual_seed(seed)
            spk_emb = chat.sample_random_speaker()
            
            # 合成音频 - 使用较低温度保持少女音特性
            params_infer_code = ChatTTS.Chat.InferCodeParams(
                spk_emb=spk_emb,
                temperature=0.2,  # 更低的温度以保持年轻音色
                top_P=0.6,
                top_K=15,
            )
            
            wavs = chat.infer([test_text], params_infer_code=params_infer_code)
            
            if wavs is not None and len(wavs) > 0:
                # 保存试听文件（主目录，方便播放）
                preview_file = f"girl_voice_{i}_{seed}.wav"
                save_wav(preview_file, wavs[0])
                
                # 保存音色文件
                speaker_file = speaker_dir / f"girl_{seed}.pt"
                torch.save(spk_emb, speaker_file)
                
                print(f"  ✓ 试听文件: {preview_file}")
                print(f"  ✓ 音色文件: {speaker_file}")
            else:
                print(f"  ✗ 生成失败")
        
        except Exception as e:
            print(f"  ✗ 错误: {e}")
    
    print("\n" + "=" * 70)
    print("生成完成!")
    print("=" * 70)
    
    print("\n📁 生成的文件:")
    for i, (seed, description) in enumerate(girl_voices.items(), 1):
        print(f"  {i}. girl_voice_{i}_{seed}.wav - {description}")
    
    print("\n" + "=" * 70)
    print("使用方法:")
    print("=" * 70)
    print("\n1. 试听上面5个音频文件，选择你喜欢的音色")
    print("\n2. 记住对应的种子号（文件名中的数字）")
    print("\n3. 编辑配置文件 config/config.yaml:")
    print("\n   tts:")
    print("     engine: \"chattts\"")
    print("     chattts:")
    print("       # 选择下面任意一种方式:")
    print("       speaker_seed: 1314  # 方式1: 使用种子号")
    print("       # 或")  
    print("       speaker_file: \"speakers/girl_1314.pt\"  # 方式2: 使用音色文件")
    print("\n4. 保存配置后运行测试:")
    print("   python test_local_enhanced.py")
    print("\n" + "=" * 70)
    print("\n💡 提示: 如果这5个都不满意，运行 python find_girl_voice.py 生成更多音色")
    print("=" * 70)


if __name__ == "__main__":
    main()
